
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <stdlib.h>
#include <stddef.h>
#include <assert.h>
#include <poll.h>
#include <errno.h>
#include <time.h>
#include <math.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <sys/uio.h>
#include <netinet/in.h>

#include "rlib.h"

#define PRINT_DEBUG 0

#define ACK_PACKET_SIZE 12
#define DATA_PACKET_HEADER_SIZE 16
#define MAX_DATA_BYTES 1000
#define MAX_WINDOW_SIZE 200

// Hack to make window size increase faster during congestion avoidance,
// leads to better empirical results
#define CONG_AVOID_INVERSE_SCALE_FACTOR 0.8
// Slides use 2.0 but anything above 1.0 slows it down because RTO too large
#define TIMEOUT_BACKOFF_MULTIPLIER 1.0
#define TIMEOUT_START 100
#define SRTT_ORIGINAL 30
#define ALPHA 1.0/4
#define BETA 1.0/8
#define GAMMA 4

static int timeouts;

typedef uint32_t seqno_t;

typedef struct send_queue_slot {
	struct timespec timeFirstSent;
	struct timespec timeLastSent;
    int use_for_rto_calc;
    int acked; // if acked then the slot is free for use
	packet_t *packet_sent;
} send_queue_slot_t;

typedef struct receive_queue_slot {
    int not_output_yet; // whether the packet has already been output or not
	packet_t *packet_received;
} receive_queue_slot_t;

struct reliable_state {

  conn_t *c;			/* This is the connection object */

  /* Add your own data fields below this */
  
	/* Sliding Window State */
    unsigned long int recv_window_size;
	int send_window_size;
	double send_windowd_size;
	int ssthresh;
	long double srtt;
	long double rttvar;
    unsigned long int timeout; // in milliseconds
	
	/* Sender state */
	seqno_t last_ack_received; // equivalent to next frame expected by receiver
	seqno_t last_frame_sent;
	int dupack_count; //Number of times Dupack received
	
	// TCP New Reno
    int fast_recovery_mode;
    seqno_t last_frame_sent_before_recover;
	
	int empty_send_queue_slots;
    send_queue_slot_t *send_queue;
	
	/* Receiver State */
	seqno_t next_frame_expected;
	receive_queue_slot_t *receive_queue;
	
	/* Flags for when to destory state */
    int read_EOF; // from input
    int received_EOF;
    int written_EOF;
    
    /* For keeping track of time to send file */
    struct timespec start_time;
	FILE *statfile;
	int bytes_sent;
};
rel_t *rel_list;

/* ---- Function Forward Declarations ------*/
void send_ack (rel_t *r);
void insert_into_receive_queue (rel_t *r, packet_t *pkt);
void handle_ack (rel_t *r, packet_t *pkt);
packet_t* new_data_pkt (char* data, uint16_t data_len, uint32_t seqno, 
    uint32_t ackno, uint32_t receiving_window_size);
unsigned long int timediff(struct timespec curr_time, struct timespec sent_time);
int time_equals(struct timespec time1, struct timespec time2);
int finished_transfer(rel_t *r);

void set_rto_karns_algorithm(rel_t* r, struct timespec sent_time);
void retransmit_packet_in_slot(rel_t* r, send_queue_slot_t *sent_slot);

/* Creates a new reliable protocol session, returns NULL on failure.
 * Exactly one of c and ss should be NULL.  (ss is NULL when called
 * from rlib.c, while c is NULL when this function is called from
 * rel_demux.) */
rel_t *
rel_create (conn_t *c, const struct sockaddr_storage *ss,
	    const struct config_common *cc)
{
  rel_t *r;

  r = xmalloc (sizeof (*r));
  memset (r, 0, sizeof (*r));

  if (!c) {
    c = conn_create (r, ss);
    if (!c) {
      free (r);
      return NULL;
    }
  }

  r->c = c;
  rel_list = r;

  /* Do any other initialization you need here */

	/* Initialize sliding window state */
    r->recv_window_size = cc->window;
	r->send_window_size = 1; 
	r->send_windowd_size = 1.0;
	r->ssthresh = 1600; // Essentially infinity
    // r->timeout = cc->timeout; // this sets it to 0 which causes tons of timeouts
    r->timeout = TIMEOUT_START;
	r->srtt = SRTT_ORIGINAL;
	r->rttvar = 0;
	
	r->last_ack_received = 1; // Other side should be expecting to receive 1 at start
	r->dupack_count = 0;
	r->last_frame_sent = 0;
	r->next_frame_expected = 1; // first frame always sent with seqno 1
	r->send_queue = (send_queue_slot_t *)malloc(sizeof(send_queue_slot_t) * MAX_WINDOW_SIZE);
    r->receive_queue = (receive_queue_slot_t *)malloc(sizeof(receive_queue_slot_t) * MAX_WINDOW_SIZE);
    
    // TCP New Reno
    r->fast_recovery_mode = 0;
    r->last_frame_sent_before_recover = 0;
    
    // Initialize queues
    size_t i;
    for (i = 0; i < MAX_WINDOW_SIZE; i++)
    {
        r->send_queue[i].packet_sent = 0;
        r->send_queue[i].timeFirstSent.tv_sec = 0;
        r->send_queue[i].timeFirstSent.tv_nsec = 0;
        r->send_queue[i].timeLastSent.tv_sec = 0; 
        r->send_queue[i].timeLastSent.tv_nsec = 0;
        r->send_queue[i].use_for_rto_calc = 0;
        r->send_queue[i].acked = 1; // since acked means slot is free
        r->receive_queue[i].packet_received = NULL;
        r->receive_queue[i].not_output_yet = 0;
    }
	
    r->read_EOF = 0;
    r->received_EOF = 0;
    r->written_EOF = 0;
    
    // Start timing
    clock_gettime(CLOCK_MONOTONIC, &(r->start_time));
	char filename[100];
	// @TODO: Filename
	sprintf(filename, "%i.stat", getpid());
	r->statfile = fopen(filename, "w");
 	if(r->statfile < 0){
		fprintf(stderr,"Error opening statistics file\n");
		sleep(1);
 	}
	r->bytes_sent = 0;
  return r;
}

void
rel_destroy (rel_t *r)
{
    conn_destroy (r->c);

    /* Free any other allocated memory here */
    struct timespec curr_time;
    clock_gettime(CLOCK_MONOTONIC, &curr_time);
    unsigned long int diff = timediff(curr_time, r->start_time);
    fprintf(stderr, "Total timeouts: %d\n", timeouts);
    fprintf(stderr, "Time taken to finish transfer: %lu ms", diff);
  
  	// free packets in sending, receiving queues then queues themselves
    size_t i;
    for (i = 0; i < MAX_WINDOW_SIZE; i++)
    {
        if (r->send_queue[i].packet_sent)
        {
            free(r->send_queue[i].packet_sent);
        }
        if (r->receive_queue[i].packet_received)
        {
            free(r->receive_queue[i].packet_received);
        }
    }
	fclose(r->statfile);
    free(r->send_queue);
    free(r->receive_queue);
}


void
rel_demux (const struct config_common *cc,
	   const struct sockaddr_storage *ss,
	   packet_t *pkt, size_t len)
{
  //leave it blank here!!!
}

void
rel_recvpkt (rel_t *r, packet_t *pkt, size_t n)
{
    uint16_t pkt_length = ntohs(pkt->len);
	
    if (n != pkt_length)
    {
		return; // packet is corrupt
    }
    
    uint16_t original_cksum = pkt->cksum;
    pkt->cksum = 0;
	if (original_cksum != cksum(pkt, n))
	{
		return; // checksum doesn't match, packet is corrupt	    
	}
	
	if (pkt_length == ACK_PACKET_SIZE) // Sending side
	{
        handle_ack(r, pkt);
	}
	else if (pkt_length >= DATA_PACKET_HEADER_SIZE) // Receiving side
	{
	    // fprintf(stderr, "received data packet, seqno %d, next frame expected %d\n", 
	    //             ntohl(pkt->seqno), r->next_frame_expected);
	    if (pkt_length == DATA_PACKET_HEADER_SIZE)
	    {
            r->received_EOF = 1;
	    }
	    
        int packet_seqno = ntohl(pkt->seqno);
        if (packet_seqno >= r->next_frame_expected &&
            packet_seqno < r->next_frame_expected + r->recv_window_size)
        {
			r->bytes_sent += pkt_length;
            insert_into_receive_queue(r, pkt);
            if (packet_seqno == r->next_frame_expected)
            {
                rel_output(r);
            }
            else // packet received out of order, send dup ACK
            {
                // fprintf(stderr, "***: sending dup ack for seqno %d\n", 
                //                     r->next_frame_expected);
                send_ack(r);
            }
        }
        else if (packet_seqno < r->next_frame_expected)
        {   // previously sent ack might have not gotten through
            send_ack(r);
        }
	}
}

void
rel_read (rel_t *s)
{
    if (s->c->sender_receiver == RECEIVER)
    {   
        // NOTE: using read_EOF as "sent EOF"
        if (s->read_EOF)
            return;
        else // send EOF
        {
            uint16_t seqno = s->last_frame_sent + 1;
            uint16_t ackno = s->next_frame_expected;
            char data[0];
            packet_t *pkt_to_send = new_data_pkt(data, 0, seqno, ackno, s->recv_window_size);
			s->bytes_sent += DATA_PACKET_HEADER_SIZE;
            conn_sendpkt(s->c, pkt_to_send, DATA_PACKET_HEADER_SIZE);
            
            // after sending, add to send queue, start timeout
            send_queue_slot_t *send_q_slot = &(s->send_queue[seqno % MAX_WINDOW_SIZE]);
            send_q_slot->packet_sent = (packet_t *)malloc(sizeof(packet_t));
            // NOTE: deep copy when inserting into send queue
            *(send_q_slot->packet_sent) = *pkt_to_send; 
            clock_gettime(CLOCK_MONOTONIC, &(send_q_slot->timeFirstSent));
            send_q_slot->timeLastSent = send_q_slot->timeFirstSent;
            send_q_slot->acked = 0;
            send_q_slot->use_for_rto_calc = 1;
            
            s->read_EOF = 1;
            s->last_frame_sent += 1;
            free(pkt_to_send);
        }
    }
    else //run in the sender mode
    {
        // if there is space, we try and send the next frame, i.e. the one after the last
        // frame sent
        // TODO: fix this to only send at most send_window_size no. of packets at a time
        //      implemented, now checking
        int next_frame_seqno = s->last_frame_sent + 1;
        while (next_frame_seqno < s->last_ack_received + s->send_window_size
               && s->send_queue[next_frame_seqno % MAX_WINDOW_SIZE].acked 
               && s->read_EOF == 0)
        {
            int input_length;
            packet_t *pkt_to_send;
            send_queue_slot_t *send_q_slot;
            char data[MAX_DATA_BYTES];

            input_length = conn_input(s->c, data, MAX_DATA_BYTES);
            if (input_length == 0) // no input
                return;
            uint16_t seqno = s->last_frame_sent + 1;
            uint16_t ackno = s->next_frame_expected;

            if (input_length == -1) // EOF or error
            {
                s->read_EOF = 1;
                pkt_to_send = new_data_pkt(data, 0, seqno, ackno, s->recv_window_size);
				s->bytes_sent += DATA_PACKET_HEADER_SIZE;
                conn_sendpkt(s->c, pkt_to_send, DATA_PACKET_HEADER_SIZE);
            }
            else
            {
                pkt_to_send = new_data_pkt(data, input_length, seqno, ackno, 
                                           s->recv_window_size);
				s->bytes_sent += input_length + DATA_PACKET_HEADER_SIZE;
                conn_sendpkt(s->c, pkt_to_send, input_length + DATA_PACKET_HEADER_SIZE);
            }

            // after sending, add to send queue, start timeout
            send_q_slot = &(s->send_queue[seqno % MAX_WINDOW_SIZE]);
            send_q_slot->packet_sent = (packet_t *)malloc(sizeof(packet_t));
            // NOTE: deep copy when inserting into send queue
            *(send_q_slot->packet_sent) = *pkt_to_send; 
            clock_gettime(CLOCK_MONOTONIC, &(send_q_slot->timeFirstSent));
            send_q_slot->timeLastSent = send_q_slot->timeFirstSent;
            send_q_slot->acked = 0;
            send_q_slot->use_for_rto_calc = 1;

            s->last_frame_sent += 1;
            next_frame_seqno = s->last_frame_sent + 1;
            // fprintf(stderr, "sending packet %d, last_ack_received %d\n", 
            //                 s->last_frame_sent, s->last_ack_received);

            free(pkt_to_send);
        }
    }
}

void
rel_output (rel_t *r)
{
    seqno_t index = r->next_frame_expected;
    while (r->receive_queue[index % MAX_WINDOW_SIZE].not_output_yet && r->written_EOF == 0)
	{
	    receive_queue_slot_t *slot = &(r->receive_queue[index % MAX_WINDOW_SIZE]);
        size_t data_size = ntohs(slot->packet_received->len) - DATA_PACKET_HEADER_SIZE;
        if (conn_bufspace(r->c) >= data_size) {
            if (data_size == 0)
            {
                r->written_EOF = 1;
            }
            conn_output(r->c, slot->packet_received->data, data_size);
            // fprintf(stderr, "+++ Output packet: %d\n", r->next_frame_expected);
        } else {
            #if PRINT_DEBUG
            // Not enough space in output buffer so send Ack of next frame to output
            fprintf(stderr, "sending ack for seqno %d\n", r->next_frame_expected);
            #endif
            send_ack(r);
            return;
        }
        slot->not_output_yet = 0;
        r->next_frame_expected += 1;
        index += 1;
	}
	#if PRINT_DEBUG
	fprintf(stderr, "sending ack for seqno %d\n", r->next_frame_expected);
	#endif
    send_ack(r);
}

void
rel_timer ()
{
    /* Retransmit any packets that need to be retransmitted */
    rel_t *r = rel_list;
    
    // fprintf(stderr, "last frame sent %d\n", r->last_frame_sent);
    struct timespec curr_time;
	clock_gettime(CLOCK_MONOTONIC, &curr_time);
	
    int num_acked = 0;
    int at_least_one_timeout = 0;
    int i;
    for (i = 0; i < r->send_window_size; i++) 
    {
        int index = (r->last_ack_received + i) % MAX_WINDOW_SIZE;
        if (r->send_queue[index].packet_sent)
        {
            if (r->send_queue[index].acked == 0)
            {	
				struct timespec sent_time = r->send_queue[index].timeLastSent;
				unsigned long int diff = timediff(curr_time, sent_time);
						
                if (diff >= r->timeout)
                {
                    #if PRINT_DEBUG
                    fprintf(stderr, "TIMEOUT OCCURED: seqno: %d, diff %lu, timeout %lu, ssthresh %i\n", ntohl(r->send_queue[index].packet_sent->seqno), diff, r->timeout, r->ssthresh);
                    #endif
                    timeouts += 1;
                    retransmit_packet_in_slot(r, &(r->send_queue[index]));
                    // Slow start because of extreme congestion
                    r->ssthresh = fmaxl(r->send_windowd_size/2, 2);
					r->send_window_size = 1;
					r->send_windowd_size = 1.0;
                    at_least_one_timeout = 1;
                }
            }
            else
            {
                num_acked += 1;
            }
        }
    }
    if (at_least_one_timeout) // doubling for every retransmitted packet is bad
    {
        // Double RTO
        r->timeout = TIMEOUT_BACKOFF_MULTIPLIER * r->timeout; // This makes it go WAY too high
    }
    
	//Print Statistics to File
	unsigned long int time = timediff(curr_time, r->start_time);
	fprintf(r->statfile, "%lu\t%f\t%i\n",time,r->send_windowd_size,r->bytes_sent*8);
	r->bytes_sent = 0;
	
    if (finished_transfer(r))
    {
        rel_destroy(r);
    }
}

/* ----------------------------
 * Helper functions
 * ----------------------------*/
 
packet_t* new_data_pkt (char* data, uint16_t data_len, uint32_t seqno, 
                        uint32_t ackno, uint32_t receiving_window_size)
{
    packet_t *pkt = malloc(sizeof(packet_t));
    memset(pkt->data, 0, MAX_DATA_BYTES); // zero out buffer
    memcpy(pkt->data, data, data_len);
    
    pkt->cksum = 0;
    pkt->seqno = htonl(seqno);
    pkt->ackno = htonl(ackno);
    pkt->rwnd = htonl(receiving_window_size);
    pkt->len = htons(data_len + DATA_PACKET_HEADER_SIZE);
    pkt->cksum = cksum(pkt, data_len + DATA_PACKET_HEADER_SIZE);
    return pkt;
}

void send_ack (rel_t* r)
{
	conn_t *c = r->c;
	seqno_t ackno = r->next_frame_expected;
	uint32_t receiving_window_size = r->recv_window_size;
    packet_t *ack = malloc(sizeof(packet_t));
    ack->cksum = 0;
    ack->len = htons(ACK_PACKET_SIZE);
    ack->ackno = htonl(ackno);
    ack->rwnd = htonl(receiving_window_size);
    ack->cksum = cksum(ack, ACK_PACKET_SIZE);
    r->bytes_sent += ACK_PACKET_SIZE;
    conn_sendpkt(c, ack, ACK_PACKET_SIZE);
    free(ack);
}

/**
 * Updates queue and last_ack_received with the latest ackno received
 */
void handle_ack (rel_t *r, packet_t *pkt)
{
    seqno_t new_ackno = ntohl(pkt->ackno);
    // fprintf(stderr, "-- ACK: new ack received ackno %d, last ack recv %d\n", new_ackno, r->last_ack_received);
    if (new_ackno > r->last_ack_received)
	{	
        // want to mark all those before new ackno as already acked
        int seqno = r->last_ack_received;
        while (seqno < new_ackno)
        {
            int index = seqno % MAX_WINDOW_SIZE;
            send_queue_slot_t *sent_slot = &(r->send_queue[index]);
            if (sent_slot->packet_sent && sent_slot->acked != 1)
            {	    
				if (sent_slot->use_for_rto_calc)
                    set_rto_karns_algorithm(r, sent_slot->timeFirstSent);
				
				if (!r->fast_recovery_mode)
				{
				    if (r->send_window_size < r->ssthresh) // Slow Start Phase
    				{
    					r->send_windowd_size += 1.0;
    					#if PRINT_DEBUG
                        fprintf(stderr, "** SLOW START: window size updated: %f\n", 
                            r->send_windowd_size);
                        #endif
    				}
    				else // Congestion avoidance phase
    				{
    					r->send_windowd_size += 1.0 / (r->send_windowd_size * 
    					    CONG_AVOID_INVERSE_SCALE_FACTOR);
    					#if PRINT_DEBUG
                        fprintf(stderr, "** CONGESTION AVOIDANCE: window size updated: %f\n", 
                            r->send_windowd_size);
                        #endif
    				}
				}
                sent_slot->acked = 1;
            }
            seqno += 1;
        }
        if (r->fast_recovery_mode)
	    {
	        if (new_ackno < r->last_frame_sent_before_recover) // partial ack
    		{
    		    retransmit_packet_in_slot(r, &(r->send_queue[new_ackno % MAX_WINDOW_SIZE]));
                int unacked = r->last_frame_sent_before_recover - new_ackno;
                double candidate_swnd = r->send_windowd_size - unacked;
                r->send_windowd_size = fmax(1.0, candidate_swnd);
                r->send_windowd_size += 1.0;
                #if PRINT_DEBUG
                fprintf(stderr, "!!! new ackno %d, recover %d\n", new_ackno, 
                    r->last_frame_sent_before_recover);
                fprintf(stderr, "** PARTIAL ACK: ack for %d, window size updated: %f\n", new_ackno,
                     r->send_windowd_size);
                #endif
    		}
    		else // full ack
    		{
                r->send_windowd_size = r->ssthresh;
                r->fast_recovery_mode = 0; // Exit fast recovery mode
                #if PRINT_DEBUG
                fprintf(stderr, "** FULL ACK: window size updated: %f\n", 
                     r->send_windowd_size);
                #endif
    		}
	    }
		r->last_ack_received = new_ackno;
		r->dupack_count = 0;
		rel_read(r); // sending window now has space
	}
	else if (new_ackno == r->last_ack_received) // Found dup ACK
	{
		r->dupack_count++;
		unsigned int retransmit_pkt_index = new_ackno % MAX_WINDOW_SIZE;
		if (r->dupack_count == 3)
		{
		    #if PRINT_DEBUG
		    fprintf(stderr, ">> 3 DUP ACK with ackno %d\n", new_ackno);
		    #endif
			// Fast retransmit for packet that received dupack for
            retransmit_packet_in_slot(r, &(r->send_queue[retransmit_pkt_index]));
            // Set all timeLastSent between dup ackno and last_frame_sent to now
            int ackno = new_ackno;
            while (ackno <= r->last_frame_sent)
            {
                struct timespec curr_time;
            	clock_gettime(CLOCK_MONOTONIC, &curr_time);
                r->send_queue[ackno % MAX_WINDOW_SIZE].timeLastSent = curr_time;
                r->send_queue[ackno % MAX_WINDOW_SIZE].use_for_rto_calc = 0;
                ackno += 1;
            }
			
			r->ssthresh = fmaxl(r->send_windowd_size/2, 2);
			r->send_windowd_size = r->ssthresh + 3;
			// TCP New Reno
            r->fast_recovery_mode = 1;
            r->last_frame_sent_before_recover = r->last_frame_sent;
		}
        else if (r->dupack_count > 3)
        {
            rel_read(r);
            r->send_windowd_size += 1.0;
        }
	}
	else if (new_ackno < r->last_ack_received)
	{
	    fprintf(stderr, "-> Out of order ack: %d, last ack recv %d\n", new_ackno, r->last_ack_received);
	}
    // r->send_window_size = floor(r->send_windowd_size);
	r->send_window_size = ceil(r->send_windowd_size);
    // fprintf(stderr, ">>> RTO: %lu\n", r->timeout);
}

void insert_into_receive_queue (rel_t *r, packet_t *pkt)
{
    seqno_t packet_seqno = ntohl(pkt->seqno);
    receive_queue_slot_t *slot = &(r->receive_queue[packet_seqno % MAX_WINDOW_SIZE]);
    slot->packet_received = (packet_t *)malloc(sizeof(packet_t));
    *(slot->packet_received) = *pkt; // Deep copy packet
    slot->not_output_yet = 1;
}

unsigned long int timediff(struct timespec curr_time, struct timespec sent_time){
	return (curr_time.tv_sec - sent_time.tv_sec) * (pow(10, 3)) 
	    + (curr_time.tv_nsec - sent_time.tv_nsec) / (pow(10, 6));
}

int time_equals(struct timespec time1, struct timespec time2){
	if(time1.tv_sec == time2.tv_sec && time1.tv_nsec == time2.tv_nsec){
		return 1;
	}
	return 0;
}

// removed r->c->sender_receiver == SENDER
int finished_transfer(rel_t *r)
{
    if (r->read_EOF && r->written_EOF 
        && r->last_ack_received > r->last_frame_sent)
        return 1;
    else
        return 0;
}

// Calculate and set RTO Karn's Algorithm
void set_rto_karns_algorithm(rel_t* r, struct timespec sent_time)
{
	struct timespec curr_time;
	clock_gettime(CLOCK_MONOTONIC, &curr_time);
	unsigned long int RTT = timediff(curr_time, sent_time);
	
	r->rttvar = BETA*abs(RTT - r->srtt) + (1-BETA)*r->rttvar;
	r->srtt = ALPHA*RTT + (1-ALPHA)*r->srtt;
	r->timeout = ceil(r->srtt + GAMMA*r->rttvar);
    // fprintf(stderr, ">>> RTT: %lu, RTO: %lu\n", RTT, r->timeout);
}

void retransmit_packet_in_slot(rel_t* r, send_queue_slot_t *sent_slot)
{
    packet_t *pkt = sent_slot->packet_sent;
	r->bytes_sent += ntohs(pkt->len);
    conn_sendpkt(r->c, pkt, ntohs(pkt->len));
    
    struct timespec curr_time;
	clock_gettime(CLOCK_MONOTONIC, &curr_time);
    sent_slot->timeLastSent = curr_time;
    sent_slot->use_for_rto_calc = 0;
}